package com.fanshuai.tree.trie.pinyin;

import com.fanshuai.tree.trie.ITrie;
import com.fanshuai.tree.trie.WordToken;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class PinyinTrie implements ITrie {
    class Node {
        public boolean isEnd;
        public Map<Character, Node> subnodes;

        public Node() {
            isEnd = false;
            subnodes = new HashMap<>();
        }

        public void addSubNode(char ch, Node node) {
            subnodes.put(ch, node);
        }
        public void delSubNode(char ch) {
            subnodes.remove(ch);
        }

        public Node getSubNode(char ch) {
            return subnodes.get(ch);
        }
    }

    private Node root;
    private int size;

    public PinyinTrie() {
        root = new Node();
        size = 0;
    }

    @Override
    public void addWord(String word) {
        if (StringUtils.isBlank(word)) {
            return;
        }

        String pinyin = PinyinUtil.getPinyinText(word);
        if (StringUtils.isBlank(pinyin)) {
            return;
        }

        Node node = root;
        for (int i = 0; i < pinyin.length(); i++) {
            char ch = pinyin.charAt(i);

            Node sub = node.getSubNode(ch);
            if (null == sub) {
                sub = new Node();
                node.addSubNode(ch, sub);
            }
            node = sub;

            if (i == pinyin.length() - 1) {
                node.isEnd = true;
                size++;
            }
        }
    }

    @Override
    public void addWords(List<String> words) {
        if (CollectionUtils.isEmpty(words)) {
            return;
        }

        words.forEach(this::addWord);
    }

    @Override
    public List<String> getAllWords() {
        List<String> allWords = new ArrayList<>(size);
        StringBuilder b = new StringBuilder();

        fillWords(root, b, allWords);
        return allWords;
    }

    private void fillWords(Node node, StringBuilder b, List<String> words) {
        if (node.isEnd) {//完整关键词
            words.add(b.toString());
        }
        if (node.subnodes.size() == 0) {//叶节点
            return;
        }

        for (Character ch : node.subnodes.keySet()) {
            b.append(ch);
            fillWords(node.getSubNode(ch), b, words);

            int lastIndex = b.toString().length() - 1;
            b.deleteCharAt(lastIndex);
        }
    }

    private WordToken pullWord(List<PinyinIndex> hit, String text) {
        String pinyin = PinyinUtil.getPinyin(hit);

        int startIndex = hit.get(0).index;
        int endIndex = hit.get(hit.size() - 1).index + 1;

        String word = text.substring(startIndex, endIndex);
        WordToken token = new WordToken(word, startIndex, endIndex);
        token.pinyin = pinyin;
        return token;
    }

    @Override
    public List<WordToken> match(String text) {
        if (StringUtils.isBlank(text) || size == 0) {
            return new ArrayList<>();
        }

        List<PinyinIndex> pinyinIndices = PinyinUtil.getPinyinIndex(text);
        if (CollectionUtils.isEmpty(pinyinIndices)) {
            return new ArrayList<>();
        }

        List<WordToken> tokens = new ArrayList<>();

        int pos = 0;
        int begin = 0;
        Node node = root;
        while (pos < pinyinIndices.size()) {
            PinyinIndex index = pinyinIndices.get(pos);
            node = node.getSubNode(index.ch);

            if (null == node) {
                begin++;
                pos = begin;
                node = root;
            } else {
                if (node.isEnd) {
                    List<PinyinIndex> hit = pinyinIndices.subList(begin, pos + 1);
                    WordToken token = pullWord(hit, text);
                    tokens.add(token);

                    pos++;
                    begin = pos;
                    node = root;
                } else {
                    pos++;
                }
            }
        }
        return tokens;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }
}
